# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import gzip

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import pandas as pd


def open_file(filepath, *args, **kwargs):
  if filepath.endswith('.gz'):
    return gzip.open(filepath, *args, **kwargs)
  else:
    return open(filepath, *args, **kwargs)


def setup_plt():
  # plt.style.use('fivethirtyeight')
  # plt.style.use('seaborn-talk')
  # plt.rcParams['figure.dpi'] = 72
  plt.rcParams['savefig.dpi'] = 100
  plt.rcParams['font.family'] = ['Inconsolata']
  plt.rcParams['font.sans-serif'] = ['Inconsolata']
  plt.rcParams['font.monospace'] = ['Inconsolata']
  # plt.rcParams['figure.constrained_layout.use'] = True

  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['xtick.labelsize'] = 7
  plt.rcParams['ytick.labelsize'] = 7


def plot_bucket(Y_pred, Y_actual, fig, gs=None, nbuckets=100):
  df = pd.DataFrame({'Y_pred': Y_pred, 'Y_actual': Y_actual})
  df.dropna(inplace=True)
  df.sort_values('Y_pred', inplace=True)

  Y_pred_mid = []
  Y_mean = []
  Y_std = []

  size = df.shape[0]
  for bucket_idx in range(nbuckets):
    from_idx = int(bucket_idx / nbuckets * size)
    to_idx = int((bucket_idx + 1) / nbuckets * size)
    Y_bucket = df.iloc[from_idx:to_idx, 1]
    Y_pred_mid.append(df.iloc[(from_idx + to_idx) // 2, 0])
    Y_mean.append(Y_bucket.mean())
    Y_std.append(Y_bucket.std())

  if gs is None:
    gs = gridspec.GridSpec(6, 1, figure=fig)
  else:
    gs = gs.subgridspec(6, 1)

  ax11 = fig.add_subplot(gs[0:2, :])
  ax11.set_xlabel('Y-pred')
  ax11.set_ylabel('Y-mean')
  ax11.plot(Y_pred_mid, Y_mean)

  ax12 = fig.add_subplot(gs[2, :])
  ax12.set_xlabel('Y-pred')
  ax12.set_ylabel('Y-std')
  ax12.plot(Y_pred_mid, Y_std)

  ax21 = fig.add_subplot(gs[3:5, :])
  ax21.set_xlabel('Y-pred %')
  ax21.set_ylabel('Y-mean')
  ax21.plot(Y_mean)

  ax22 = fig.add_subplot(gs[5, :])
  ax22.set_xlabel('Y-pred %')
  ax22.set_ylabel('Y-std')
  ax22.plot(Y_std)
